#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
============================================================
JSON / JSONL 基础与本脚本功能说明
============================================================
1. JSON 是什么？
   - 全称 JavaScript Object Notation，一种轻量级数据交换格式。
   - 基本数据类型：对象(object, 用花括号 { } 包含键值对)、数组(array, 用方括号 [ ])、字符串、数字、布尔、null。
   - JSON 中的键必须是双引号包裹的字符串，字符串内部用反斜杠转义特殊字符（如 "\n"）。
   - Python 使用内置模块 json：json.loads(str) 解析字符串，json.dumps(obj, ensure_ascii=False) 序列化对象。

2. JSONL (JSON Lines) 是什么？
   - 也称 NDJSON，每一行都是一个独立的完整 JSON 对象，行与行之间用换行符分隔。
   - 适用场景：
	   * 流式/增量处理（可以逐行读取，内存占用小）。
	   * 日志、批量推理结果、大规模数据集。
   - 与“一个大 JSON 数组”相比优点：
	   * 无需一次性加载整个文件。
	   * 出错定位简单（按行号定位）。
	   * 可以方便地追加（append）新行。
   - 读取方式：逐行读取 -> 去掉空行 / 注释 -> json.loads(该行)。

3. 本数据集的特殊结构说明：
   - 每一行（一个 JSON 对象）包含字段：
	   * custom_id: 例如 "01-03-0001#2_W"
	   * model:     例如 "qwen2.5-vl-72b-instruct"
	   * response:  一个“数组(list) 或 字符串”，当前文件中是 list，内部含有一个 Markdown 代码块形式的 JSON 字符串，例如：
		   [ "```json\n{\n  \"target_epoch_stage\": \"W\", ... }\n```" ]
   - 目标：
	   * 解析外层 JSON 行。
	   * 从 response 中提取真正的 JSON 内容（去掉 ```json 代码围栏）。
	   * 解析得到内部字段：target_epoch_stage, reasoning_text, applicable_rules。

4. Markdown 代码块与嵌套 JSON：
   - 由于内部 JSON 被包裹在三反引号代码块 ```json ... ``` 中，需要先用正则或字符串裁剪去掉围栏。
   - 处理策略：优先用正则匹配三反引号；若未匹配到，则尝试直接 json.loads。
   - 可能的异常：
	   * 代码块缺失或不完整 -> 记录警告。
	   * 字符串里包含多余的注释或非 JSON 内容 -> 清理后再尝试。

5. 错误处理 / 鲁棒性建议：
   - 行级别解析失败不要终止整个程序，打印到 stderr 并继续。
   - 允许通过命令行参数设置最大解析行数 (--limit)。
   - 可选择输出为 CSV / JSONL 以便后续处理。

6. 本脚本提供内容：
   - parse_jsonl_line: 解析单行并返回结构化数据。
   - extract_inner_json: 从 response 字段中抽出真正 JSON 字符串并解析。
   - main: 命令行接口，支持：
	   * 输入文件 (--input)
	   * 输出文件 (--output, 可选)
	   * 输出格式 (--format csv|jsonl)
	   * 限制条目数 (--limit N)
	   * 显示进度 (--show-progress)

7. 快速使用示例：
   - 仅打印前 5 条：
	   python parse_jsonl.py --input ../MASS-SS3_Qwen2.5-VL-72B-Instruct/01-03-0001.jsonl --limit 5
   - 保存为 CSV：
	   python parse_jsonl.py -i ../MASS-SS3_Qwen2.5-VL-72B-Instruct/01-03-0001.jsonl -o parsed.csv -f csv
   - 保存为 JSONL：
	   python parse_jsonl.py -i ../MASS-SS3_Qwen2.5-VL-72B-Instruct/01-03-0001.jsonl -o parsed.jsonl -f jsonl

============================================================
"""

from __future__ import annotations

import argparse
import csv
import json
import re
import sys
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union


# ------------------------------ 数据结构定义 ------------------------------
@dataclass
class InnerResponse:
	"""表示嵌套 JSON（去除 Markdown 代码块）后解析出的字段。"""

	target_epoch_stage: Optional[str] = None
	reasoning_text: Optional[str] = None
	applicable_rules: Optional[str] = None


@dataclass
class ParsedLine:
	"""表示单行 JSONL 解析后的规范结构。"""

	line_number: int
	custom_id: Optional[str]
	model: Optional[str]
	target_epoch_stage: Optional[str]
	reasoning_text: Optional[str]
	applicable_rules: Optional[str]
	raw_response: Optional[str]  # 原始的（去围栏前）response 字符串，便于调试


# ------------------------------ 工具函数 ------------------------------
CODE_BLOCK_REGEX = re.compile(
	r"```(?:json)?\n(?P<body>[\s\S]*?)```", re.IGNORECASE
)


def normalize_newlines(s: str) -> str:
	"""统一换行符为 \n。"""

	return s.replace("\r\n", "\n").replace("\r", "\n")


def extract_inner_json(response_obj: Any) -> Tuple[str, Dict[str, Any]]:
	"""从 response 字段中提取并解析内层 JSON。

	参数:
		response_obj: 可能是 str、list[str]、其它类型。

	返回:
		(raw_text, parsed_dict)
		raw_text: 去除首尾空白后的原始（或拼接后）字符串（含或不含代码围栏）。
		parsed_dict: 解析成功的 dict；若失败，返回 {}。
	"""

	# 1. 统一成字符串
	if isinstance(response_obj, list):
		# 经验：该数据集中 response 是长度为 1 的列表
		raw = "".join(str(part) for part in response_obj)
	else:
		raw = str(response_obj)

	raw = normalize_newlines(raw).strip()

	# 2. 正则匹配 Markdown 代码块
	match = CODE_BLOCK_REGEX.search(raw)
	inner_text = match.group("body").strip() if match else raw

	# 3. 尝试直接 json.loads
	parsed: Dict[str, Any] = {}
	try:
		parsed = json.loads(inner_text)
	except json.JSONDecodeError:
		# 有时内部可能带有多余 ``` 或者缩进问题，尝试二次清理
		cleaned = inner_text.strip()
		# 去掉可能残留的反引号
		cleaned = cleaned.strip("`")
		try:
			parsed = json.loads(cleaned)
		except json.JSONDecodeError as e2:
			# 解析失败，返回空 dict，同时 stderr 提示
			# print(
			# 	f"[WARN] 内层 JSON 解析失败: {e2} | 内容片段: {cleaned[:80]!r}",
			# 	file=sys.stderr,
			# )
			parsed = {}

	return raw, parsed


def parse_jsonl_line(line: str, line_number: int) -> Optional[ParsedLine]:
	"""解析单个 JSONL 行。

	失败策略：
		- 如果外层 json 无法解析，返回 None。
		- 内层 response 解析失败：对应字段设为 None。
	"""

	text = line.strip()
	if not text:
		return None
	try:
		outer = json.loads(text)
	except json.JSONDecodeError as e:
		print(f"[ERROR] 第 {line_number} 行外层 JSON 解析失败: {e}", file=sys.stderr)
		return None

	custom_id = outer.get("custom_id")
	model = outer.get("model")
	response_obj = outer.get("response")

	raw_response, inner = extract_inner_json(response_obj)
	inner_resp = InnerResponse(
		target_epoch_stage=inner.get("target_epoch_stage"),
		reasoning_text=inner.get("reasoning_text"),
		applicable_rules=inner.get("applicable_rules"),
	) if inner else InnerResponse()

	return ParsedLine(
		line_number=line_number,
		custom_id=custom_id,
		model=model,
		target_epoch_stage=inner_resp.target_epoch_stage,
		reasoning_text=inner_resp.reasoning_text,
		applicable_rules=inner_resp.applicable_rules,
		raw_response=raw_response,
	)


def iter_jsonl(path: Union[str, Path]) -> Iterable[str]:
    path = Path(path)  # 确保转成 Path 对象
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            yield line


# ------------------------------ 输出函数 ------------------------------
def write_csv(rows: List[ParsedLine], output: Path) -> None:
	"""将结果写入 CSV。"""

	fieldnames = [
		"line_number",
		"custom_id",
		"model",
		"target_epoch_stage",
		"applicable_rules",
		"reasoning_text",
	]
	# reasoning_text 放最后，防止换行影响浏览器/表格预览
	with output.open("w", encoding="utf-8", newline="") as f:
		writer = csv.DictWriter(f, fieldnames=fieldnames)
		writer.writeheader()
		for r in rows:
			row_dict = {
				"line_number": r.line_number,
				"custom_id": r.custom_id,
				"model": r.model,
				"target_epoch_stage": r.target_epoch_stage,
				"applicable_rules": r.applicable_rules,
				# 将换行替换为 \n 方便在 Excel 单元格中查看；也可去掉替换保持原样
				"reasoning_text": (r.reasoning_text or "").replace("\n", "\\n"),
			}
			writer.writerow(row_dict)


def write_jsonl(rows: List[ParsedLine], output: Path) -> None:
	"""写入 JSONL：每行一个纯 JSON。"""

	with output.open("w", encoding="utf-8") as f:
		for r in rows:
			obj = asdict(r)
			# raw_response 可能很长，如果不需要可注释下一行
			# 也可以：del obj['raw_response']
			json_str = json.dumps(obj, ensure_ascii=False)
			f.write(json_str + "\n")


# ------------------------------ 命令行接口 ------------------------------
def build_arg_parser() -> argparse.ArgumentParser:
	p = argparse.ArgumentParser(
		description="解析含嵌套 JSON 的 JSONL 文件，提取自定义字段。"
	)
	p.add_argument(
		"-i",
		"--input",
		type=Path,
		required=True,
		help="输入 JSONL 文件路径",
	)
	p.add_argument(
		"-o",
		"--output",
		type=Path,
		default=None,
		help="输出文件路径（可选，不提供则只打印前若干条）",
	)
	p.add_argument(
		"-f",
		"--format",
		choices=["csv", "jsonl"],
		default="csv",
		help="输出格式：csv 或 jsonl (默认 csv)",
	)
	p.add_argument(
		"--limit",
		type=int,
		default=None,
		help="最多解析多少条（用于快速预览、调试）",
	)
	p.add_argument(
		"--show-progress",
		action="store_true",
		help="是否在 stderr 打印进度 (每 1000 行)",
	)
	return p


def main(argv: Optional[List[str]] = None) -> int:
	args = build_arg_parser().parse_args(argv)
	if not args.input.exists():
		print(f"[ERROR] 输入文件不存在: {args.input}", file=sys.stderr)
		return 1

	rows: List[ParsedLine] = []
	for idx, line in enumerate(iter_jsonl(args.input), start=1):
		parsed = parse_jsonl_line(line, idx)
		if parsed:
			rows.append(parsed)
		if args.show_progress and idx % 1000 == 0:
			print(f"[INFO] 已处理 {idx} 行", file=sys.stderr)
		if args.limit is not None and len(rows) >= args.limit:
			break

	# 如果未指定输出文件，只打印一个简要预览
	if args.output is None:
		preview_n = min(5, len(rows))
		print(f"共解析 {len(rows)} 条，预览前 {preview_n} 条：")
		for r in rows[:preview_n]:
			print(
				f"line={r.line_number} custom_id={r.custom_id} model={r.model} "
				f"stage={r.target_epoch_stage} rules={r.applicable_rules}\n"
				f"reasoning={ (r.reasoning_text or '')[:120] + ('...' if r.reasoning_text and len(r.reasoning_text)>120 else '') }"
			)
		print("(使用 -o 保存完整结果；--format 选择输出格式)\n")
		return 0

	# 写入输出
	if args.format == "csv":
		write_csv(rows, args.output)
	else:
		write_jsonl(rows, args.output)
	print(f"[OK] 写入完成: {args.output} ({len(rows)} 条)")
	return 0


if __name__ == "__main__":  # 允许作为脚本运行
	sys.exit(main())

